{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4-5 AutoGraph和tf.Module"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "有三种计算图的构建方式：静态计算图，动态计算图，以及Autograph。\n",
    "\n",
    "TensorFlow 2.0主要使用的是动态计算图和Autograph。\n",
    "\n",
    "* 动态计算图易于调试，编码效率较高，但执行效率偏低。\n",
    "\n",
    "* 静态计算图执行效率很高，但较难调试。\n",
    "\n",
    "* Autograph机制可以将动态图转换成静态计算图，兼收执行效率和编码效率之利。\n",
    "\n",
    "当然Autograph机制能够转换的代码并不是没有任何约束的，有一些编码规范需要遵循，否则可能会转换失败或者不符合预期。\n",
    "\n",
    "前面我们介绍了Autograph的编码规范和Autograph转换成静态图的原理。\n",
    "\n",
    "本篇我们介绍使用tf.Module来更好地构建Autograph。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 一、Autograph和tf.Module概述\n",
    "\n",
    "前面在介绍Autograph的编码规范时提到构建Autograph时应该避免在@tf.function修饰的函数内部定义tf.Variable. \n",
    "\n",
    "但是如果在函数外部定义tf.Variable的话，又会显得这个函数有外部变量依赖，封装不够完美。\n",
    "\n",
    "一种简单的思路是定义一个类，并将相关的tf.Variable创建放在类的初始化方法中。而将函数的逻辑放在其他方法中。\n",
    "\n",
    "这样一顿猛如虎的操作之后，我们会觉得一切都如同人法地地法天天法道道法自然般的自然。\n",
    "\n",
    "惊喜的是，TensorFlow提供了一个基类tf.Module，通过继承它构建子类，我们不仅可以获得以上的自然而然，而且可以非常方便地管理变量，还可以非常方便地管理它引用的其它Module，最重要的是，我们能够利用tf.saved_model保存模型并实现跨平台部署使用。\n",
    "\n",
    "实际上，tf.keras.models.Model，tf.keras.layers.Layer 都是继承自tf.Module的，提供了方便的变量管理和所引用的子模块管理的功能。\n",
    "\n",
    "**因此，利用tf.Module提供的封装，再结合TensoFlow丰富的低阶API，实际上我们能够基于TensorFlow开发任意机器学习模型(而非仅仅是神经网络模型)，并实现跨平台部署使用。**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 二、应用tf.Module封装Autograph\n",
    "\n",
    "定义一个简单的function。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "x = tf.Variable(1.0, dtype=tf.float32)\n",
    "\n",
    "import tensorflow as tf \n",
    "x = tf.Variable(1.0,dtype=tf.float32)\n",
    "\n",
    "#在tf.function中用input_signature限定输入张量的签名类型：shape和dtype\n",
    "@tf.function(input_signature=[tf.TensorSpec(shape = [], dtype = tf.float32)])    \n",
    "def add_print(a):\n",
    "    x.assign_add(a)\n",
    "    tf.print(x)\n",
    "    return(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=float32, numpy=4.0>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "add_print(tf.constant(3.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Python inputs incompatible with input_signature:\n  inputs: (\n    tf.Tensor(3, shape=(), dtype=int32))\n  input_signature: (\n    TensorSpec(shape=(), dtype=tf.float32, name=None))",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-4-1c1afb3b3caa>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0madd_print\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconstant\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m#输入不符合张量签名的参数将报错\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwds)\u001b[0m\n\u001b[1;32m    566\u001b[0m         \u001b[0mxla_context\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mExit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    567\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 568\u001b[0;31m       \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    569\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    570\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mtracing_count\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_tracing_count\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py\u001b[0m in \u001b[0;36m_call\u001b[0;34m(self, *args, **kwds)\u001b[0m\n\u001b[1;32m    604\u001b[0m       \u001b[0;31m# In this case we have not created variables on the first call. So we can\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    605\u001b[0m       \u001b[0;31m# run the first trace but we should fail if variables are created.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 606\u001b[0;31m       \u001b[0mresults\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_stateful_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    607\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_created_variables\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    608\u001b[0m         raise ValueError(\"Creating variables on a non-first call to a function\"\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   2360\u001b[0m     \u001b[0;34m\"\"\"Calls a graph function specialized to the inputs.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2361\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_lock\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2362\u001b[0;31m       \u001b[0mgraph_function\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_maybe_define_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2363\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mgraph_function\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_filtered_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2364\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py\u001b[0m in \u001b[0;36m_maybe_define_function\u001b[0;34m(self, args, kwargs)\u001b[0m\n\u001b[1;32m   2659\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minput_signature\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0margs\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mkwargs\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2660\u001b[0m       args, kwargs = self._function_spec.canonicalize_function_inputs(\n\u001b[0;32m-> 2661\u001b[0;31m           *args, **kwargs)\n\u001b[0m\u001b[1;32m   2662\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2663\u001b[0m     \u001b[0mcache_key\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_cache_key\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py\u001b[0m in \u001b[0;36mcanonicalize_function_inputs\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   2183\u001b[0m           \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2184\u001b[0m           \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_input_signature\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2185\u001b[0;31m           self._flat_input_signature)\n\u001b[0m\u001b[1;32m   2186\u001b[0m       \u001b[0;32mreturn\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2187\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py\u001b[0m in \u001b[0;36m_convert_inputs_to_signature\u001b[0;34m(inputs, input_signature, flat_input_signature)\u001b[0m\n\u001b[1;32m   2250\u001b[0m       flatten_inputs)):\n\u001b[1;32m   2251\u001b[0m     raise ValueError(\"Python inputs incompatible with input_signature:\\n%s\" %\n\u001b[0;32m-> 2252\u001b[0;31m                      format_error_message(inputs, input_signature))\n\u001b[0m\u001b[1;32m   2253\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2254\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0mneed_packing\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: Python inputs incompatible with input_signature:\n  inputs: (\n    tf.Tensor(3, shape=(), dtype=int32))\n  input_signature: (\n    TensorSpec(shape=(), dtype=tf.float32, name=None))"
     ]
    }
   ],
   "source": [
    "add_print(tf.constant(3)) #输入不符合张量签名的参数将报错"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "下面利用tf.Module的子类化将其封装一下。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DemoModule(tf.Module):\n",
    "    def __init__(self, init_value=tf.constant(0.0), name=None):\n",
    "        super(DemoModule, self).__init__(name=name)\n",
    "        with self.name_scope: # 相当于with tf.name_scope(\"demo_module\")\n",
    "            self.x = tf.Variable(init_value, dtype=tf.float32, trainable=True)\n",
    "            \n",
    "    @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.float32)])\n",
    "    def addprint(self, a):\n",
    "        with self.name_scope:\n",
    "            self.x.assign_add(a)\n",
    "            tf.print(self.x)\n",
    "            return self.x\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6\n"
     ]
    }
   ],
   "source": [
    "# 执行\n",
    "demo = DemoModule(init_value=tf.constant(1.0))\n",
    "result = demo.addprint(tf.constant(5.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(<tf.Variable 'demo_module/Variable:0' shape=() dtype=float32, numpy=6.0>,)\n",
      "(<tf.Variable 'demo_module/Variable:0' shape=() dtype=float32, numpy=6.0>,)\n"
     ]
    }
   ],
   "source": [
    "#查看模块中的全部变量和全部可训练变量\n",
    "print(demo.variables)\n",
    "print(demo.trainable_variables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "()"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#查看模块中的全部子模块\n",
    "demo.submodules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/alan/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/resource_variable_ops.py:1786: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n",
      "INFO:tensorflow:Assets written to: ./data/demo/1/assets\n"
     ]
    }
   ],
   "source": [
    "#使用tf.saved_model 保存模型，并指定需要跨平台部署的方法\n",
    "tf.saved_model.save(demo, \"./data/demo/1\", signatures={\"serving_default\":demo.addprint})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=float32, numpy=11.0>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 加载模型\n",
    "demo2 = tf.saved_model.load(\"./data/demo/1\")\n",
    "demo2.addprint(tf.constant(5.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:\n",
      "\n",
      "signature_def['__saved_model_init_op']:\n",
      "  The given SavedModel SignatureDef contains the following input(s):\n",
      "  The given SavedModel SignatureDef contains the following output(s):\n",
      "    outputs['__saved_model_init_op'] tensor_info:\n",
      "        dtype: DT_INVALID\n",
      "        shape: unknown_rank\n",
      "        name: NoOp\n",
      "  Method name is: \n",
      "\n",
      "signature_def['serving_default']:\n",
      "  The given SavedModel SignatureDef contains the following input(s):\n",
      "    inputs['a'] tensor_info:\n",
      "        dtype: DT_FLOAT\n",
      "        shape: ()\n",
      "        name: serving_default_a:0\n",
      "  The given SavedModel SignatureDef contains the following output(s):\n",
      "    outputs['output_0'] tensor_info:\n",
      "        dtype: DT_FLOAT\n",
      "        shape: ()\n",
      "        name: StatefulPartitionedCall:0\n",
      "  Method name is: tensorflow/serving/predict\n",
      "WARNING:tensorflow:From /Users/alan/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/resource_variable_ops.py:1786: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n",
      "\n",
      "Defined Functions:\n",
      "  Function Name: 'addprint'\n",
      "    Option #1\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          a: TensorSpec(shape=(), dtype=tf.float32, name='a')\n"
     ]
    }
   ],
   "source": [
    "# 查看模型文件相关信息，红框标出来的输出信息在模型部署和跨平台使用时有可能会用到\n",
    "!saved_model_cli show --dir ./data/demo/1 --all"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](./data/查看模型文件信息.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在tensorboard中查看计算图，模块会被添加模块名demo_module，方便层次化呈现计算图结构。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5\n"
     ]
    }
   ],
   "source": [
    "import datetime\n",
    "\n",
    "# 创建日志\n",
    "stamp = datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "logdir = \"./data/demomodule/%s\" % stamp\n",
    "writer = tf.summary.create_file_writer(logdir)\n",
    "\n",
    "# 开启 autograph跟踪\n",
    "tf.summary.trace_on(graph=True, profiler=True)\n",
    "\n",
    "# 执行 autograph\n",
    "demo = DemoModule(init_value=tf.constant(0.0))\n",
    "result = demo.addprint(tf.constant(5.0))\n",
    "\n",
    "# 将计算图信息写入日志 \n",
    "with writer.as_default():\n",
    "    tf.summary.trace_export(\n",
    "        name=\"demomodule\",\n",
    "        step=0,\n",
    "        profiler_outdir=logdir\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "#启动 tensorboard在jupyter中的魔法命令\n",
    "%reload_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Known TensorBoard instances:\n",
      "  - port 6006: logdir ./data/keras_model (started 1 day, 0:49:19 ago; pid 23825)\n",
      "  - port 6007: logdir ./data/autograph/ (started 17:58:28 ago; pid 29609)\n"
     ]
    }
   ],
   "source": [
    "from tensorboard import notebook\n",
    "notebook.list() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "      <iframe id=\"tensorboard-frame-e34eed02bc6da95b\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
       "      </iframe>\n",
       "      <script>\n",
       "        (function() {\n",
       "          const frame = document.getElementById(\"tensorboard-frame-e34eed02bc6da95b\");\n",
       "          const url = new URL(\"/\", window.location);\n",
       "          url.port = 6008;\n",
       "          frame.src = url;\n",
       "        })();\n",
       "      </script>\n",
       "  "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "notebook.start(\"--logdir ./data/demomodule/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "除了利用tf.Module的子类化实现封装，我们也可以通过给tf.Module添加属性的方法进行封装。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "mymodule = tf.Module()\n",
    "mymodule.x = tf.Variable(0.0)\n",
    "\n",
    "@tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.float32)])\n",
    "def addprint(a):\n",
    "    mymodule.x.assign_add(a)\n",
    "    tf.print(mymodule.x)\n",
    "    return mymodule.x\n",
    "\n",
    "mymodule.addprint = addprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mymodule.addprint(tf.constant(1.0)).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=3.0>,)\n"
     ]
    }
   ],
   "source": [
    "print(mymodule.variables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: data/mymodule/assets\n",
      "8\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=float32, numpy=8.0>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 使用 tf.saved_model保存模型\n",
    "tf.saved_model.save(mymodule, \"data/mymodule\", \n",
    "                   signatures = {\"serving_default\": mymodule.addprint})\n",
    "\n",
    "# 加载模型\n",
    "mymodule2 = tf.saved_model.load(\"./data/mymodule\")\n",
    "mymodule2.addprint(tf.constant(5.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 三、tf.Module和tf.keras.Model，tf.keras.layers.Layer\n",
    "\n",
    "tf.keras中的模型和层都是继承tf.Module实现的，也具有变量管理和子模块管理功能。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.keras import models, layers, losses, metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "print(issubclass(tf.keras.Model, tf.Module))\n",
    "print(issubclass(tf.keras.layers.Layer, tf.Module))\n",
    "print(issubclass(tf.keras.Model, tf.Module))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "dense (Dense)                (None, 4)                 44        \n",
      "_________________________________________________________________\n",
      "dense_1 (Dense)              (None, 2)                 10        \n",
      "_________________________________________________________________\n",
      "dense_2 (Dense)              (None, 1)                 3         \n",
      "=================================================================\n",
      "Total params: 57\n",
      "Trainable params: 57\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "tf.keras.backend.clear_session()\n",
    "\n",
    "model = models.Sequential()\n",
    "\n",
    "model.add(layers.Dense(4, input_shape=(10,)))\n",
    "model.add(layers.Dense(2))\n",
    "model.add(layers.Dense(1))\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<tf.Variable 'dense/kernel:0' shape=(10, 4) dtype=float32, numpy=\n",
       " array([[ 0.40079212, -0.115282  ,  0.2583183 , -0.61419123],\n",
       "        [ 0.48109078, -0.6408005 ,  0.5820811 , -0.03453928],\n",
       "        [-0.43774536, -0.55664444,  0.3507135 ,  0.38044858],\n",
       "        [-0.6068879 , -0.54923534, -0.49395925, -0.6472933 ],\n",
       "        [-0.48002052,  0.23096955, -0.18679425, -0.34857765],\n",
       "        [ 0.4077741 ,  0.5153465 ,  0.04699057, -0.16955972],\n",
       "        [-0.4933897 , -0.19523919, -0.3902719 , -0.5296963 ],\n",
       "        [-0.25946817,  0.11002082, -0.45470035,  0.21353537],\n",
       "        [ 0.05975348, -0.1832323 , -0.05176556,  0.37643373],\n",
       "        [ 0.40585542, -0.00723654, -0.42971325,  0.5353552 ]],\n",
       "       dtype=float32)>,\n",
       " <tf.Variable 'dense/bias:0' shape=(4,) dtype=float32, numpy=array([0., 0., 0., 0.], dtype=float32)>,\n",
       " <tf.Variable 'dense_1/kernel:0' shape=(4, 2) dtype=float32, numpy=\n",
       " array([[-0.8640754 ,  0.62785053],\n",
       "        [-0.4134264 ,  0.18177843],\n",
       "        [ 0.25406456, -0.73353744],\n",
       "        [ 0.14676237,  0.93448997]], dtype=float32)>,\n",
       " <tf.Variable 'dense_1/bias:0' shape=(2,) dtype=float32, numpy=array([0., 0.], dtype=float32)>,\n",
       " <tf.Variable 'dense_2/kernel:0' shape=(2, 1) dtype=float32, numpy=\n",
       " array([[ 1.0810379],\n",
       "        [-1.3192958]], dtype=float32)>,\n",
       " <tf.Variable 'dense_2/bias:0' shape=(1,) dtype=float32, numpy=array([0.], dtype=float32)>]"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<tf.Variable 'dense_1/kernel:0' shape=(4, 2) dtype=float32, numpy=\n",
       " array([[-0.8640754 ,  0.62785053],\n",
       "        [-0.4134264 ,  0.18177843],\n",
       "        [ 0.25406456, -0.73353744],\n",
       "        [ 0.14676237,  0.93448997]], dtype=float32)>,\n",
       " <tf.Variable 'dense_1/bias:0' shape=(2,) dtype=float32, numpy=array([0., 0.], dtype=float32)>,\n",
       " <tf.Variable 'dense_2/kernel:0' shape=(2, 1) dtype=float32, numpy=\n",
       " array([[ 1.0810379],\n",
       "        [-1.3192958]], dtype=float32)>,\n",
       " <tf.Variable 'dense_2/bias:0' shape=(1,) dtype=float32, numpy=array([0.], dtype=float32)>]"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.layers[0].trainable = False #冻结第 0 层的变量， 使其不可训练\n",
    "model.trainable_variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<tensorflow.python.keras.engine.input_layer.InputLayer at 0x14823c080>,\n",
       " <tensorflow.python.keras.layers.core.Dense at 0x14820c358>,\n",
       " <tensorflow.python.keras.layers.core.Dense at 0x14823c5f8>,\n",
       " <tensorflow.python.keras.layers.core.Dense at 0x14823cb70>)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.submodules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<tensorflow.python.keras.layers.core.Dense at 0x14820c358>,\n",
       " <tensorflow.python.keras.layers.core.Dense at 0x14823c5f8>,\n",
       " <tensorflow.python.keras.layers.core.Dense at 0x14823cb70>]"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sequential\n",
      "sequential\n"
     ]
    }
   ],
   "source": [
    "print(model.name)\n",
    "print(model.name_scope())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
